#include "ecies.h"

namespace umtl
{
	class ecies_tester
	{
	public:
		int test();

	private:
		void processor_cleanup(unsigned char *text, unsigned char *copy, unsigned char *original);
		int processor(int iteration);
	};

	void ecies_tester::processor_cleanup(unsigned char *text, unsigned char *copy, unsigned char *original)
	{
		if (text) {
			free(text);
		}

		if (copy) {
			free(copy);
		}

		if (original) {
			free(original);
		}

		return;
	}

	int ecies_tester::processor(int iteration)
	{
		int tlen;
		size_t olen;
		ecies::cryptex_t ciphered;
		char* hex_pub, *hex_priv;
		unsigned char *text = NULL, *copy = NULL, *original = NULL;

		// Generate random size for the block of data were going to encrypt. Use a min value of 1 KB and a max of 1 MB.
		do {
			tlen = (rand() % (1024 * 1024));
		} while (tlen < 1024);

		if (!(text = (unsigned char *)malloc(tlen + 1)) || !(copy = (unsigned char *)malloc(tlen + 1))) {
			printf("Memory error.\n");
			processor_cleanup(text, copy, original);
			return -1;
		}

		// Wipe and then fill the data blocks with random data.
		memset(copy, 0, tlen + 1);
		memset(text, 0, tlen + 1);
		for (uint64_t j = 0; j < tlen; j++) {
			*(copy + j) = *(text + j) = (rand() % 255);
		}

		ecies::ec_key_t& ec_key = ecies::create_keys();

		// Generate a key for our theoretical user.
		if (!(ec_key.get())) {
			printf("Key creation failed.\n");
			processor_cleanup(text, copy, original);
			return -1;
		}


		// Since we'll store the keys as hex values in reali life, extract the appropriate hex values and release the original key structure.
		if (!(hex_pub = ecies::public_key()) || !(hex_priv = ecies::private_key())) {
				printf("Serialization of the key to a pair of hex strings failed.\n");
				processor_cleanup(text, copy, original);
				return -1;
		}
		
		if (!(ciphered = ecies::Instance().ecies_encrypt(hex_pub, text, tlen))) {
			printf("The encryption process failed!\n");
			processor_cleanup(text, copy, original);
			return -1;
		}

		if (!(original = ecies::Instance().ecies_decrypt(hex_priv, ciphered, &olen))) {
			printf("The decryption process failed!\n");
			processor_cleanup(text, copy, original);
			return -1;
		}

		if (olen != tlen || memcmp(original, copy, tlen)) {
			printf("Comparison failure.\n");
			processor_cleanup(text, copy, original);
			return -1;
		}

		processor_cleanup(text, copy, original);
		printf(" ... %i ... %i\n", iteration + 1, tlen);

		return 0;
	}

	int ecies_tester::test() {

		// Comment this line out if you want the program to execute consistently each time.
		srand(time(NULL));

		for (uint64_t i = 0; i < 100; i++) {
			if (processor(i)) {
				return 1;
			}
		}

		printf("Finished.\n");

		return 0;
	}
}